{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:\n",
      "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "  * https://github.com/tensorflow/io (for I/O related ops)\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n",
      "WARNING:tensorflow:From /home/sudharsan/anaconda3/envs/universe/lib/python3.6/site-packages/tensorflow/contrib/layers/python/layers/layers.py:1634: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.flatten instead.\n",
      "WARNING:tensorflow:Entity <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd4081f6eb8>> could not be transformed and will be executed as-is. Please report this to the AutgoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: converting <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd4081f6eb8>>: AttributeError: module 'gast' has no attribute 'Num'\n",
      "WARNING: Entity <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd4081f6eb8>> could not be transformed and will be executed as-is. Please report this to the AutgoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: converting <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd4081f6eb8>>: AttributeError: module 'gast' has no attribute 'Num'\n",
      "WARNING:tensorflow:Entity <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd3c1e2a5f8>> could not be transformed and will be executed as-is. Please report this to the AutgoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: converting <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd3c1e2a5f8>>: AttributeError: module 'gast' has no attribute 'Num'\n",
      "WARNING: Entity <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd3c1e2a5f8>> could not be transformed and will be executed as-is. Please report this to the AutgoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: converting <bound method Flatten.call of <tensorflow.python.layers.core.Flatten object at 0x7fd3c1e2a5f8>>: AttributeError: module 'gast' has no attribute 'Num'\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import random\n",
    "from collections import deque\n",
    "import math\n",
    "\n",
    "import gym\n",
    "from tensorflow.python.framework import ops\n",
    "\n",
    "\n",
    "def conv(inputs, kernel_shape, bias_shape, strides, weights, bias=None, activation=tf.nn.relu):\n",
    "\n",
    "    weights = tf.get_variable('weights', shape=kernel_shape, initializer=weights)\n",
    "    conv = tf.nn.conv2d(inputs, weights, strides=strides, padding='SAME')\n",
    "    if bias_shape is not None:\n",
    "        biases = tf.get_variable('biases', shape=bias_shape, initializer=bias)\n",
    "        return activation(conv + biases) if activation is not None else conv + biases\n",
    "    return activation(conv) if activation is not None else conv\n",
    "\n",
    "\n",
    "def dense(inputs, units, bias_shape, weights, bias=None, activation=tf.nn.relu):\n",
    "    \n",
    "\n",
    "    if not isinstance(inputs, ops.Tensor):\n",
    "        inputs = ops.convert_to_tensor(inputs, dtype='float')\n",
    "    if len(inputs.shape) > 2:\n",
    "        inputs = tf.contrib.layers.flatten(inputs)\n",
    "    flatten_shape = inputs.shape[1]\n",
    "    weights = tf.get_variable('weights', shape=[flatten_shape, units], initializer=weights)\n",
    "    dense = tf.matmul(inputs, weights)\n",
    "    if bias_shape is not None:\n",
    "        assert bias_shape[0] == units\n",
    "        biases = tf.get_variable('biases', shape=bias_shape, initializer=bias)\n",
    "        return activation(dense + biases) if activation is not None else dense + biases\n",
    "    return activation(dense) if activation is not None else dense\n",
    "\n",
    "v_min = 0\n",
    "v_max = 1000\n",
    "atoms = 51\n",
    "gamma = 0.99 \n",
    "batch_size = 10\n",
    "update_target_net = 50  \n",
    "epsilon = 0.5\n",
    "\n",
    "buffer_length = 20000\n",
    "replay_buffer = deque(maxlen=buffer_length)\n",
    "\n",
    "def sample_transitions(batch_size):\n",
    "    batch = np.random.permutation(len(replay_buffer))[:batch_size]\n",
    "    trans = np.array(replay_buffer)[batch]\n",
    "    return trans\n",
    "\n",
    "\n",
    "\n",
    "class Categorical_DQN():\n",
    "    def __init__(self,env):\n",
    "        self.sess = tf.InteractiveSession()\n",
    "        self.v_max = v_max\n",
    "        self.v_min = v_min\n",
    "        self.atoms = atoms \n",
    "\n",
    "        self.epsilon = epsilon\n",
    "        self.state_shape = env.observation_space.shape\n",
    "        self.action_shape = env.action_space.n\n",
    "\n",
    "        self.time_step = 0\n",
    "\n",
    "        target_state_shape = [1]\n",
    "        target_state_shape.extend(self.state_shape)\n",
    "\n",
    "\n",
    "        self.state_ph = tf.placeholder(tf.float32,target_state_shape)\n",
    "        self.action_ph = tf.placeholder(tf.int32,[1,1])\n",
    "        self.m_ph = tf.placeholder(tf.float32,[self.atoms])\n",
    "\n",
    "        self.delta_z = (self.v_max - self.v_min) / (self.atoms - 1)\n",
    "        self.z = [self.v_min + i * self.delta_z for i in range(self.atoms)]\n",
    "\n",
    "        self.build_categorical_DQN()\n",
    "        self.sess.run(tf.global_variables_initializer())\n",
    "\n",
    "\n",
    "    def build_network(self, state, action, name, units_1, units_2, weights, bias, reg=None):\n",
    "        with tf.variable_scope('conv1'):\n",
    "            conv1 = conv(state, [5, 5, 3, 6], [6], [1, 2, 2, 1], weights, bias)\n",
    "        with tf.variable_scope('conv2'):\n",
    "            conv2 = conv(conv1, [3, 3, 6, 12], [12], [1, 2, 2, 1], weights, bias)\n",
    "        with tf.variable_scope('flatten'):\n",
    "            flatten = tf.contrib.layers.flatten(conv2)\n",
    "\n",
    "        with tf.variable_scope('dense1'):\n",
    "            dense1 = dense(flatten, units_1, [units_1], weights, bias)\n",
    "        with tf.variable_scope('dense2'):\n",
    "            dense2 = dense(dense1, units_2, [units_2], weights, bias)\n",
    "        with tf.variable_scope('concat'):\n",
    "            concatenated = tf.concat([dense2, tf.cast(action, tf.float32)], 1)\n",
    "        with tf.variable_scope('dense3'):\n",
    "            dense3 = dense(concatenated, self.atoms, [self.atoms], weights, bias) \n",
    "        return tf.nn.softmax(dense3)\n",
    "\n",
    "    def build_categorical_DQN(self):\n",
    "        with tf.variable_scope('target_net'):\n",
    "            name = ['target_net_params',tf.GraphKeys.GLOBAL_VARIABLES]\n",
    "\n",
    "            weights = tf.random_uniform_initializer(-0.1,0.1)\n",
    "            bias = tf.constant_initializer(0.1)\n",
    "\n",
    "            self.target_p = self.build_network(self.state_ph,self.action_ph,name,24,24,weights,bias)\n",
    "\n",
    "        with tf.variable_scope('main_net'):\n",
    "            name = ['main_net_params',tf.GraphKeys.GLOBAL_VARIABLES]\n",
    "            weights = tf.random_uniform_initializer(-0.1,0.1)\n",
    "            bias = tf.constant_initializer(0.1)\n",
    "\n",
    "            self.main_p = self.build_network(self.state_ph,self.action_ph,name,24,24,weights,bias)\n",
    "\n",
    "\n",
    "        self.main_Q = tf.reduce_sum(self.main_p * self.z)\n",
    "        self.target_Q = tf.reduce_sum(self.target_p * self.z)\n",
    "\n",
    "        self.cross_entropy_loss = -tf.reduce_sum(self.m_ph * tf.log(self.main_p))\n",
    "        self.optimizer = tf.train.AdamOptimizer(0.01).minimize(self.cross_entropy_loss)\n",
    "\n",
    "        main_net_params = tf.get_collection(\"main_net_params\")\n",
    "        target_net_params = tf.get_collection('target_net_params')\n",
    "\n",
    "        self.update_target_net = [tf.assign(t, e) for t, e in zip(target_net_params, main_net_params)]\n",
    "\n",
    "\n",
    "    def train(self,s,r,action,s_,gamma):\n",
    "        self.time_step += 1\n",
    "\n",
    "        list_q_ = [self.sess.run(self.target_Q,feed_dict={self.state_ph:[s_],self.action_ph:[[a]]}) for a in range(self.action_shape)]\n",
    "        \n",
    "        a_ = tf.argmax(list_q_).eval()\n",
    "        \n",
    "\n",
    "        m = np.zeros(self.atoms)\n",
    "        p = self.sess.run(self.target_p,feed_dict = {self.state_ph:[s_],self.action_ph:[[a_]]})[0]\n",
    "        for j in range(self.atoms):\n",
    "            Tz = min(self.v_max,max(self.v_min,r+gamma * self.z[j]))\n",
    "            bj = (Tz - self.v_min) / self.delta_z \n",
    "            l,u = math.floor(bj),math.ceil(bj) \n",
    "\n",
    "            pj = p[j]\n",
    "\n",
    "            m[int(l)] += pj * (u - bj)\n",
    "            m[int(u)] += pj * (bj - l)\n",
    "\n",
    "        self.sess.run(self.optimizer,feed_dict={self.state_ph:[s] , self.action_ph:[action], self.m_ph: m })\n",
    "        if self.time_step % update_target_net == 0:\n",
    "            self.sess.run(self.update_target_net)\n",
    "\n",
    "    def select_action(self,s):\n",
    "        if random.random() <= self.epsilon:\n",
    "            return random.randint(0, self.action_shape - 1)\n",
    "        else: \n",
    "            return np.argmax([self.sess.run(self.main_Q,feed_dict={self.state_ph:[s],self.action_ph:[[a]]}) for a in range(self.action_shape)])\n",
    "\n",
    "num_episodes = 800\n",
    "env = gym.make(\"Tennis-v0\")\n",
    "agent = Categorical_DQN(env)\n",
    "\n",
    "\n",
    "for i in range(num_episodes):\n",
    "    done = False\n",
    "    state = env.reset()\n",
    "    Return = 0\n",
    "\n",
    "    while not done:\n",
    "\n",
    "        env.render()\n",
    "        action = agent.select_action(state)\n",
    "\n",
    "        next_state, reward, done, info = env.step(action)\n",
    "\n",
    "        Return = Return + reward\n",
    "\n",
    "        replay_buffer.append([state, reward, [action], next_state])\n",
    "\n",
    "        if len(replay_buffer) >= batch_size:\n",
    "            trans = sample_transitions(2)\n",
    "            for item in trans:\n",
    "                agent.train(item[0],item[1], item[2], item[3],gamma)\n",
    "\n",
    "        state = next_state\n",
    "            \n",
    "    print(\"Episode:{}, Return: {}\".format(i,Return))\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
